from asyncio import FastChildWatcher
from code import interact
import os
import time
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
#from isaacgym import gymtorch
from math import sqrt
import math
from sympy import false
import torch
import cv2

from draw import *

from pcgworker.PCGWorker import *

from wfc_env import *

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.results_plotter import load_results, ts2xy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common import results_plotter
from stable_baselines3.common.results_plotter import load_results, ts2xy, plot_results
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.callbacks import EvalCallback, StopTrainingOnRewardThreshold

class SaveOnBestTrainingRewardCallback(BaseCallback):
    """
    Callback for saving a model (the check is done every ``check_freq`` steps)
    based on the training reward (in practice, we recommend using ``EvalCallback``).

    :param check_freq:
    :param log_dir: Path to the folder where the model will be saved.
      It must contains the file created by the ``Monitor`` wrapper.
    :param verbose: Verbosity level.
    """
    def __init__(self, check_freq: int, log_dir: str, verbose: int = 1, model_name: str="best_model", stop_reward = 0.82, active_steps = 10000):
        super(SaveOnBestTrainingRewardCallback, self).__init__(verbose)
        self.check_freq = check_freq
        self.log_dir = log_dir
        self.save_path = os.path.join(log_dir, model_name)
        self.best_mean_reward = -np.inf

        # start timer
        self.start = time.time()

        self.model_name = model_name
        self.stop_reward = stop_reward
        self.active_steps = active_steps

    def _init_callback(self) -> None:
        # Create folder if needed
        if self.save_path is not None:
            os.makedirs(self.save_path, exist_ok=True)
        self.start = time.time()

    def _on_step(self) -> bool:
        if self.n_calls % self.check_freq == 0:
            end = time.time()
            
            # Retrieve training reward
            x, y = ts2xy(load_results(self.log_dir), 'timesteps')
            if len(x) > 0:
                # Mean training reward over the last 100 episodes
                mean_reward = np.mean(y[-100:])
                # if self.verbose > 0:
                #     print("Time taken: {}".format(end - self.start))
                #     print(f"Num timesteps: {self.num_timesteps}")
                #     print(f"Best mean reward: {self.best_mean_reward:.2f} - Last mean reward per episode: {mean_reward:.2f}")

                # New best model, you could save the agent here
                if mean_reward > self.best_mean_reward and self.num_timesteps > self.active_steps:
                    self.best_mean_reward = mean_reward
                    # Example for saving best model
                    if self.verbose > 0:
                        print(f"Saving new best model to {self.save_path}")
                        self.model.save(self.save_path)
                        
                    # stop training procedure by modifiying num_timesteps of the model
                    if self.best_mean_reward >= self.stop_reward:
                        self.model.num_timesteps = 10000000

        return True

LOGDIR = "./training_logs"
timesteps = 600000

# Create the callback: check every 1000 steps
callback = SaveOnBestTrainingRewardCallback(check_freq=1000, log_dir=LOGDIR, verbose=1, model_name="best_model", stop_reward = 0.91)

m_env = Monitor(CustomEnv( render_indicator=False ), LOGDIR, allow_early_resets=True)
observation = m_env.reset()

policy_kwargs = dict(activation_fn=torch.nn.ReLU,
                     net_arch=[dict(pi=[256, 256], vf=[256, 256])])

# model = PPO('CnnPolicy', m_env, policy_kwargs=policy_kwargs).learn(total_timesteps=timesteps,callback=callback)
model = PPO('CnnPolicy', m_env)

# stage 1 
model.learn(total_timesteps=timesteps,callback=callback)
# model.learn(total_timesteps=timesteps,callback=eval_callback)

# stage 2
callback1 = SaveOnBestTrainingRewardCallback(check_freq=1000, log_dir=LOGDIR, verbose=1, model_name="best_model10", stop_reward = 0.82)
m_env.close()
m_env1 = Monitor(CustomEnv(headless_ = True, seed_pth = "seed.json", render_indicator=False), LOGDIR, allow_early_resets=True)
m_env1.reset()

model.set_env(m_env1,force_reset = False)
model.learn(total_timesteps=(timesteps*4),callback=callback1)

print("best mean reward: =================================================== ", callback1.best_mean_reward)